import numpy as np
import matplotlib.pyplot as plt
from limsim import TrafficSimulator, BEVVisualizer


class MultiScenarioSimulator:
    def __init__(self):
        # 初始化仿真器（集成LimSim环岛与高速路网）
        self.sim = TrafficSimulator(scenario="roundabout_and_highway")
        self.bev_viz = BEVVisualizer(resolution=(800, 800))

        # 车辆行为参数
        self.aggressive_prob = 0.2
        self.conservative_prob = 0.3
        self.safe_time_gap = 2.0  # 安全时间间隔（秒）

    def set_vehicle_behavior(self, vehicle_id):
        """设置异质化驾驶行为模型"""
        if np.random.rand() < self.aggressive_prob:
            # 激进型：高加速度，低安全裕度
            return {"accel_range": [1.5, 2.5], "min_gap": 0.8}
        elif np.random.rand() < self.conservative_prob:
            # 保守型：低加速度，高安全裕度
            return {"accel_range": [-1.0, 1.0], "min_gap": 2.5}
        else:
            # 正常型：中等参数
            return {"accel_range": [0.5, 1.5], "min_gap": 1.5}

    def run_roundabout_scenario(self):
        """环岛多车交互逻辑"""
        ego_id = self.sim.add_vehicle(type="car", start_lane="entrance_1", target_exit=3)
        for _ in range(15):  # 添加背景车辆
            veh_id = self.sim.add_vehicle(type="car",
                                          start_lane=np.random.choice(["entrance_1", "entrance_2", "entrance_3"]))
            self.sim.set_driving_profile(veh_id, self.set_vehicle_behavior(veh_id))

        for step in range(200):  # 仿真200步（20秒）
            # Ego车辆决策：基于环岛内车辆TTC（Time-to-Collision）
            island_vehicles = self.sim.get_vehicles_in_roundabout()
            if island_vehicles:
                closest_veh = min(island_vehicles, key=lambda v: self.sim.calc_ttc(ego_id, v))
                if self.sim.calc_ttc(ego_id, closest_veh) < self.safe_time_gap:
                    self.sim.set_acceleration(ego_id, -1.0)  # 减速让行
                else:
                    self.sim.set_acceleration(ego_id, 1.5)  # 加速进入
            self.sim.step()

    def run_highway_overtaking(self):
        """高速变道超车逻辑"""
        ego_id = self.sim.add_vehicle(type="car", start_lane="left", velocity=90)
        truck_id = self.sim.add_vehicle(type="truck", start_lane="middle", velocity=80)
        for _ in range(8):  # 添加随机轿车
            veh_id = self.sim.add_vehicle(type="car", start_lane=np.random.choice(["middle", "right"]))
            self.sim.set_driving_profile(veh_id, {"accel_range": [-1.5, 1.5]})

        for step in range(150):  # 仿真150步（15秒）
            # Ego车辆决策：右侧车道空间检测
            right_lane_free = self.sim.check_lane_availability(ego_id, "right", gap_threshold=20)
            if self.sim.get_vehicle_ahead(ego_id, same_lane=True).speed < 85 and right_lane_free:
                self.sim.change_lane(ego_id, "right")  # 变道超车
            self.sim.step()

    def visualize_bev(self):
        """生成动态鸟瞰图并可视化"""
        bev_image = self.bev_viz.render(
            vehicles=self.sim.get_vehicle_states(),
            road_network=self.sim.get_road_geometry(),
            traffic_lights=self.sim.get_traffic_light_states()
        )
        plt.imshow(bev_image)
        plt.title("BEV: Roundabout & Highway Overtaking")
        plt.axis('off')
        plt.show()


if __name__ == "__main__":
    simulator = MultiScenarioSimulator()
    simulator.run_roundabout_scenario()
    simulator.run_highway_overtaking()
    simulator.visualize_bev()